package websocket

import (
	"sync"

	"github.com/gorilla/websocket"
)

var (
	MessageList = 10
	GoWrite     = 2
)

type WebSockerWriteGroup struct {
	messageType int

	connMutex *sync.Mutex
	conns     map[*websocket.Conn]struct{}

	messageList chan []byte
	// goWSize     chan struct{} // 并发写到socket 的数量

	delConn []*websocket.Conn

	isClose bool
}

func NewWebSockerWriteGroup(messageType int) *WebSockerWriteGroup {
	ws := &WebSockerWriteGroup{
		messageType: messageType,

		connMutex: new(sync.Mutex),
		conns:     make(map[*websocket.Conn]struct{}),

		messageList: make(chan []byte, MessageList),
		// goWSize:     make(chan struct{}, GoWrite),
		delConn: make([]*websocket.Conn, 0, 2),
	}
	go ws.run()
	return ws
}

func (pool *WebSockerWriteGroup) run() {
	var data []byte
	var err error
	for {
		if pool.isClose {
			break
		}

		data = <-pool.messageList
		pool.connMutex.Lock()
		for conn, _ := range pool.conns {
			err = conn.WriteMessage(pool.messageType, data)
			if err != nil {
				pool.delConn = append(pool.delConn, conn)
			}
		}

		for _, conn := range pool.delConn {
			delete(pool.conns, conn)
		}
		pool.connMutex.Unlock()

	}
}

func (pool *WebSockerWriteGroup) Write(data []byte) {
	pool.messageList <- data
	return
}

func (pool *WebSockerWriteGroup) AddConn(conn *websocket.Conn) {
	pool.connMutex.Lock()
	pool.conns[conn] = struct{}{}
	pool.connMutex.Unlock()
}

func (pool *WebSockerWriteGroup) DelConn(conn *websocket.Conn) {
	pool.delConn = append(pool.delConn, conn)
}

func (pool *WebSockerWriteGroup) Close() {
	pool.isClose = true
	pool.conns = nil
	pool.Write([]byte("Close"))
}
